# import nibabel as nib
import numpy as np
import os
import SimpleITK as sitk
import matplotlib.pyplot as plt
from skimage import data, draw, color, transform, feature,measure
import random

import cv2
import numpy as np
import math

import cv2 as cv
import numpy as np
import copy
from PIL import Image

from skimage.morphology import erosion, dilation

def multi_dilation(image, kernel, iterations):
    for i in range(iterations):
        image = dilation(image, kernel)
    return image

def connected_component(image):
    # 标记输入的3D图像
    label, num = measure.label(image, connectivity=1, return_num=True)
    # print(label)
    # print(num)
    if num < 1:
        return image
    else:
        cross = np.array([[0, 1, 0],
                          [1, 1, 1],
                          [0, 1, 0]])
        image = multi_dilation(image, cross, 1)
        return image

def connected_domain_2(image, mask=True):
    cca = sitk.ConnectedComponentImageFilter()
    cca.SetFullyConnected(True)
    _input = sitk.GetImageFromArray(image.astype(np.uint8))
    output_ex = cca.Execute(_input)
    stats = sitk.LabelShapeStatisticsImageFilter()
    stats.Execute(output_ex)
    num_label = cca.GetObjectCount()
    num_list = [i for i in range(1, num_label+1)]
    area_list = []
    for l in range(1, num_label +1):
        area_list.append(stats.GetNumberOfPixels(l))
    num_list_sorted = sorted(num_list, key=lambda x: area_list[x-1])[::-1]
    largest_area = area_list[num_list_sorted[0] - 1]
    final_label_list = [num_list_sorted[0]]

    for idx, i in enumerate(num_list_sorted[1:]):
        if area_list[i-1] >= (largest_area//10):
            final_label_list.append(i)
        else:
            break
    output = sitk.GetArrayFromImage(output_ex)

    for one_label in num_list:
        if  one_label in final_label_list:
            continue
        x, y, z, w, h, d = stats.GetBoundingBox(one_label)
        one_mask = (output[z: z + d, y: y + h, x: x + w] != one_label)
        output[z: z + d, y: y + h, x: x + w] *= one_mask

    if mask:
        output = (output > 0).astype(np.uint8)
    else:
        output = ((output > 0)*255.).astype(np.uint8)
    return output


def post_process(in_path, out_path):
    image_array = sitk.GetArrayFromImage(sitk.ReadImage(in_path))
    image_array = image_array.astype(np.uint8)
    out_array = connected_domain_2(image_array)
    sitk.WriteImage(sitk.GetImageFromArray(out_array), out_path)
    


# files = os.listdir('D:/code/FFR-heart/prediction_177-188/')
# for file in files:
#     image_path = os.path.join('D:/code/FFR-heart/prediction_177-188/',file)
#     labelImage = sitk.ReadImage(image_path)  # in_file是nii.gz文件的路径
#     imag = sitk.GetArrayFromImage(labelImage)
#     vol = imag.astype('uint8')
#     print(image_path)
#     vol = connected_domain_2(vol)

#     # get image data
#     image_out = sitk.GetImageFromArray(vol)
#     sitk.WriteImage(image_out, os.path.join('D:/code/FFR-heart/after_post_process',file))
